#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from rclpy.clock import Clock
import random
import socket
import time
import string
import argparse
from threading import Lock
from std_msgs.msg import Header
from agv_msgs.msg import AgvAutoStatus, AGVEvent, AGVMode

class AGVSimulator(Node):
    def __init__(self, ip, port, tcp_interval_ms):
        super().__init__('agv_simulator')
        
        # 将毫秒转换为秒
        tcp_interval_sec = max(0.001, tcp_interval_ms / 1000.0)  # 最小1毫秒
        
        self.get_logger().info(
            f"AGV Simulator Started. TCP Server: {ip}:{port}, "
            f"Interval: {tcp_interval_ms}ms ({tcp_interval_sec:.3f}s)"
        )
        
        # TCP 客户端参数
        self.tcp_ip = ip
        self.tcp_port = port
        self.tcp_interval_ms = tcp_interval_ms
        self.tcp_interval_sec = tcp_interval_sec
        self.tcp_socket = None
        self.tcp_connected = False
        self.tcp_lock = Lock()  # 用于线程安全的TCP操作
        
        # 创建状态发布者
        self.status_pub = self.create_publisher(AgvAutoStatus, '/auto_run_status', 10)
        self.status_timer = self.create_timer(1.0, self.simulate_agv_status)  # 1Hz
        
        # 创建事件发布者
        self.event_pub = self.create_publisher(AGVEvent, '/eventstatus', 10)
        self.event_timer = self.create_timer(5.0, self.simulate_event)  # 0.2Hz
        
        # 创建模式发布者
        self.mode_pub = self.create_publisher(AGVMode, '/agvmode', 10)
        self.mode_timer = self.create_timer(10.0, self.simulate_agv_mode)  # 0.1Hz
        
        # 创建TCP通信定时器（使用毫秒转换后的秒值）
        self.tcp_timer = self.create_timer(tcp_interval_sec, self.tcp_communication)
        
        # 预定义事件代码
        self.event_codes = [
            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
            21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
            40, 41, 42, 50, 52, 53, 54, 60, 61, 62, 63, 70, 100, 101, 102, 103, 110, 
            111, 200, 201, 202, 203, 204, 205, 206
        ]
    
    def connect_to_server(self):
        """连接到TCP服务器"""
        try:
            with self.tcp_lock:
                if self.tcp_socket:
                    self.tcp_socket.close()
                
                self.tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.tcp_socket.settimeout(max(0.01, self.tcp_interval_sec * 0.5))  # 设置合理超时
                self.tcp_socket.connect((self.tcp_ip, self.tcp_port))
                self.tcp_connected = True
                self.get_logger().info(f"Connected to TCP server at {self.tcp_ip}:{self.tcp_port}")
                return True
        except Exception as e:
            self.get_logger().error(f"TCP connection failed: {str(e)}")
            self.tcp_connected = False
            return False
    
    def tcp_communication(self):
        """处理TCP通信：发送和接收数据"""
        # 确保连接状态
        if not self.tcp_connected:
            if not self.connect_to_server():
                self.get_logger().warning("TCP connection not established. Retrying...")
                return
        
        try:
            start_time = time.time()  # 记录开始时间用于性能监控
            with self.tcp_lock:
                # 生成发送数据：前4字节为"WD: "，后12字节为随机数据
                prefix = b"WD: "
                data = ''.join(random.choices(
                    string.ascii_letters + string.digits + string.punctuation, 
                    k=12
                )).encode('utf-8')
                
                message = prefix + data
                self.get_logger().debug(f"Generated data: {data.decode('utf-8', errors='replace')}")
                
                # 发送数据
                self.tcp_socket.sendall(message)
                self.get_logger().debug(f"Sent TCP data: {message.hex()}")
                
               # 接收数据 (16字节) - 改进版本
                response = b''
                remaining = 16
                recv_timeout = 5.0  # 总接收超时时间
                start_recv = time.time()
                
                while remaining > 0:
                    try:
                        # 设置每次recv的超时
                        self.tcp_socket.settimeout(1.0)  # 每次recv等待1秒
                        
                        chunk = self.tcp_socket.recv(remaining)
                        if not chunk:
                            self.get_logger().warning("Connection closed by server during receive")
                            return
                            
                        response += chunk
                        remaining -= len(chunk)
                        self.get_logger().debug(f"Received chunk: {len(chunk)} bytes, total: {len(response)}/16")
                        
                        # 检查总接收时间
                        if time.time() - start_recv > recv_timeout:
                            self.get_logger().warning(f"Receive timed out after {recv_timeout}s, got {len(response)}/16 bytes")
                            return
                            
                    except socket.timeout:
                        elapsed = time.time() - start_recv
                        self.get_logger().warning(f"Partial receive timeout: {elapsed:.2f}s, received {len(response)}/16 bytes")
                        # 不立即返回，继续尝试直到总超时
                        continue
                        
                self.get_logger().debug(f"Full response: {response.hex()}")
                
                # 验证响应数据
                if len(response) < 5:
                    self.get_logger().error("Response too short for validation")
                    return
                    
                # 检查ACK前缀
                if response[:5] != b"ACK: ":
                    self.get_logger().error(f"Invalid response prefix. Expected 'ACK: ', got {response[:5]}")
                    return
                    
                # 检查数据前缀匹配
                if len(response) >= 9 and response[5:9] != message[4:8]:
                    self.get_logger().error(f"Data mismatch. Sent prefix: {message[4:8]}, Received: {response[5:9]}")
                    return
                    
                # 计算通信耗时
                duration_ms = (time.time() - start_time) * 1000
                
                # 所有验证通过
                self.get_logger().info(
                    f"TCP communication validated successfully! "
                    f"Duration: {duration_ms:.2f}ms"
                )
                self.get_logger().debug(f"Received data matches sent prefix: {response[5:9].decode()}")
                
        except socket.timeout:
            self.get_logger().warning("TCP communication timed out")
            # self.tcp_connected = False
        except socket.error as e:
            self.get_logger().error(f"TCP communication error: {str(e)}")
            self.tcp_connected = False
        except Exception as e:
            self.get_logger().error(f"Unexpected error in TCP communication: {str(e)}")
            # self.tcp_connected = False
    
    def simulate_agv_status(self):
        msg = AgvAutoStatus()
        clock = Clock()
        now = clock.now().to_msg()
        
        # 设置header
        msg.header = Header()
        msg.header.stamp = now
        msg.header.frame_id = "base_link"
        
        # 填充数据
        msg.taskid = random.randint(1, 10000)
        msg.subtaskid = random.randint(1, 100)
        msg.fromnode = random.randint(1, 100)
        msg.tonode = random.randint(1, 100)
        msg.fprogress = random.uniform(0.0, 100.0)
        msg.fprogpercent = random.uniform(0.0, 100.0)
        msg.curpst_x = random.uniform(-10.0, 10.0)
        msg.curpst_y = random.uniform(-10.0, 10.0)
        msg.curpst_t = random.uniform(-3.14, 3.14)
        msg.curvel_x = random.uniform(0.0, 2.0)
        msg.curvel_z = random.uniform(-1.0, 1.0)
        msg.vellevel = random.randint(0, 3)
        msg.fdevprog = random.uniform(-0.5, 0.5)
        msg.fdev = random.uniform(-0.3, 0.3)
        msg.fdevthita = random.uniform(-0.2, 0.2)
        msg.fquality = random.uniform(0.5, 1.0)
        msg.event = random.randint(0, 100)
        msg.workstate = random.choice([0, 1, 2])
        msg.wheelmovedir = random.choice([0, 1])
        msg.trajtype = random.choice([-1, 0, 1, 2, 3])
        msg.locatetype = random.choice([0, 1, 2])
        msg.scenes = "scene_" + str(random.randint(1, 5))
        msg.bits = random.randint(0, 65535)
        msg.word1 = random.randint(-100, 100)
        msg.word2 = random.randint(-100, 100)
        msg.pathword1 = random.randint(-100, 100)
        msg.pathword2 = random.randint(-100, 100)
        msg.rackfollowup = random.choice([0, 1])
        msg.navigationtype = random.choice([0, 1, 2])
        msg.driving = random.choice([True, False])
        msg.startdelay = random.choice([True, False])
        msg.paused = random.choice([True, False])
        msg.idle = random.choice([True, False])
        msg.orderidnew = "order_" + str(random.randint(1000, 9999))
        msg.orderupdateidnew = random.randint(1, 1000)
        msg.nodeidnew = "node_" + str(random.randint(1, 50))
        
        self.status_pub.publish(msg)
        self.get_logger().debug('Published AGV Status', throttle_duration_sec=5.0)
    
    def simulate_event(self):
        msg = AGVEvent()
        clock = Clock()
        now = clock.now().to_msg()
        
        # 设置header
        msg.header = Header()
        msg.header.stamp = now
        msg.header.frame_id = "base_link"
        
        # 填充数据
        msg.type = random.choice([0, 1, 2])
        msg.action = random.choice([0, 1, 2, 3, 4])
        msg.ackreq = random.choice([0, 1])
        msg.eventtime = now.sec + now.nanosec / 1e9
        msg.vellimit = random.uniform(0.1, 2.0)
        msg.ecode = random.choice(self.event_codes)
        msg.param1 = random.randint(-100, 100)
        msg.param2 = random.randint(-100, 100)
        msg.param3 = random.randint(-100, 100)
        msg.param4 = random.randint(-100, 100)
        msg.description_cn = f"事件描述_{random.randint(1, 100)}"
        msg.description = f"Event Description_{random.randint(1, 100)}"
        
        self.event_pub.publish(msg)
        self.get_logger().info('Published Event')
    
    def simulate_agv_mode(self):
        msg = AGVMode()
        clock = Clock()
        now = clock.now().to_msg()
        
        # 设置header
        msg.header = Header()
        msg.header.stamp = now
        msg.header.frame_id = "base_link"
        
        # 填充数据
        msg.mode = random.choice([0, 1, 2, 3])
        
        self.mode_pub.publish(msg)
        self.get_logger().info(f'Published AGV Mode: {msg.mode}')

def main(args=None):
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='AGV Simulator with TCP Client')
    parser.add_argument('--ip', type=str, default='127.0.0.1', 
                        help='TCP server IP address (default: 127.0.0.1)')
    parser.add_argument('--port', type=int, default=50001, 
                        help='TCP server port (default: 50001)')
    parser.add_argument('--tcp_interval', type=float, default=100.0,
                        help='TCP communication interval in milliseconds (default: 100ms)')
    
    # 解析ROS2参数和自定义参数
    args, _ = parser.parse_known_args()
    
    # 确保最小间隔为1ms
    tcp_interval_ms = max(1.0, args.tcp_interval)
    
    rclpy.init(args=None)
    simulator = AGVSimulator(
        ip=args.ip, 
        port=args.port, 
        tcp_interval_ms=tcp_interval_ms
    )
    
    try:
        rclpy.spin(simulator)
    except KeyboardInterrupt:
        pass
    finally:
        # 清理TCP连接
        with simulator.tcp_lock:
            if simulator.tcp_socket:
                simulator.tcp_socket.close()
                simulator.get_logger().info("TCP connection closed")
        
        simulator.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()